{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "import torch\n",
    "# 测试 取出百川模型的embedding权重，并重新加载\n",
    "llama_model = AutoModelForCausalLM.from_pretrained(\n",
    "                \"/home/41_data/models--baichuan-inc--Baichuan2-7B-Chat/snapshots/ea66ced17780ca3db39bc9f8aa601d8463db3da5\",\n",
    "                # torch_dtype=torch.float32 if is_inference else torch.float16,\n",
    "                torch_dtype=torch.float32,\n",
    "                trust_remote_code=True,\n",
    "                output_hidden_states=True,\n",
    "            )\n",
    "\n",
    "embedding_model = llama_model.model.embed_tokens\n",
    "weight = embedding_model.weight # (12w, 4096)\n",
    "weight_path = \"./data_list/gxl_embedding.pt\"\n",
    "torch.save(weight, weight_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "from gxl_ai_utils.utils import utils_file\n",
    "utils_file.hello_gxl()"
   ],
   "metadata": {
    "collapsed": false,
    "is_executing": true
   },
   "id": "e0c6137a4a21dbaf"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "print('llama_model')"
   ],
   "metadata": {
    "collapsed": false,
    "is_executing": true
   },
   "id": "5c4780696f251c8a"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "embedding_weights = torch.load('embedding_weights.pt')\n",
    "embedding_model = torch.nn.Embedding.from_pretrained(embedding_weights)\n",
    "input_ids = torch.randint(4,12000,size=(4,92))\n",
    "with torch.no_grad():\n",
    "    output = embedding_model(input_ids)\n",
    "print(output.shape)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "80d08604f13fd017"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import torch \n",
    "text2embedding = torch.nn.Embedding.from_pretrained(\n",
    "            torch.load(\"/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/gxl_embedding.pt\"))\n",
    "input_inx = torch.tensor([12,32,23,45343,23,23,-100,-100])\n",
    "output_inx = text2embedding(input_inx)\n",
    "print(output_inx.shape)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "97e0958769e6cea1"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
